﻿namespace Jsl.FxCop
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using Microsoft.FxCop.Sdk;

    /// <summary>
    /// Builds the statement flow control for a <see cref="Method"/>.
    /// </summary>
    public class StatementFlow
    {
        private readonly Method method;

        private Dictionary<Statement, StatementInfo> statements;

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the <see cref="StatementFlow"/> class.
        /// </summary>
        /// <param name="method">The method to get statement flow from.</param>
        public StatementFlow(Method method)
        {
            if (method == null)
            {
                throw new ArgumentNullException("method");
            }

            this.method = method;
        }
        #endregion Constructors

        #region Public Properties
        /// <summary>
        /// Gets the method.
        /// </summary>
        /// <value>The method.</value>
        public Method Method
        {
            get { return this.method; }
        }

        /// <summary>
        /// Gets the all the statements in the method.
        /// </summary>
        /// <value>All the statements.</value>
        public IEnumerable<Statement> Statements
        {
            get
            {
                this.EnsureFlowHasBeenBuilt();

                return this.statements.Keys;
            }
        }
        #endregion Public Properties

        #region Public Methods
        /// <summary>
        /// Returns all that paths that start with the specified statement.
        /// </summary>
        /// <param name="initialStatement">The initial statement.</param>
        /// <returns>All the paths that start with the specified statement.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1006:DoNotNestGenericTypesInMemberSignatures", Justification = "By design.")]
        public IEnumerable<IEnumerable<Statement>> AllFollowingPaths(Statement initialStatement)
        {
            return this.AllPaths(initialStatement, this.FollowingStatements);
        }

        /// <summary>
        /// Returns all that paths that end with the specified statement.
        /// </summary>
        /// <param name="initialStatement">The initial statement.</param>
        /// <returns>All the paths that end with the specified statement.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1006:DoNotNestGenericTypesInMemberSignatures", Justification = "By design.")]
        public IEnumerable<IEnumerable<Statement>> AllPreviousPaths(Statement initialStatement)
        {
            return this.AllPaths(initialStatement, this.PreviousStatements);
        }

        /// <summary>
        /// Enumerates all the statements that follow this statement.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>All the statements that follow this statement.</returns>
        public IEnumerable<Statement> AllFollowingStatements(Statement statement)
        {
            return this.AllStatements(statement, this.FollowingStatements);
        }

        /// <summary>
        /// Enumerates all the statements previous to this statement.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>All the statements previous to this statement.</returns>
        public IEnumerable<Statement> AllPreviousStatements(Statement statement)
        {
            return this.AllStatements(statement, this.PreviousStatements);
        }

        /// <summary>
        /// Gets the following statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The following statements, if any.</returns>
        public IEnumerable<Statement> FollowingStatements(Statement statement)
        {
            this.EnsureFlowHasBeenBuilt();

            StatementInfo statementInfo;
            if (this.statements.TryGetValue(statement, out statementInfo))
            {
                return statementInfo.Following;
            }
            else
            {
                return Enumerable.Empty<Statement>();
            }
        }

        /// <summary>
        /// Gets the previous statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The previous statements, if any.</returns>
        public IEnumerable<Statement> PreviousStatements(Statement statement)
        {
            this.EnsureFlowHasBeenBuilt();

            StatementInfo statementInfo;
            if (this.statements.TryGetValue(statement, out statementInfo))
            {
                return statementInfo.Previous;
            }
            else
            {
                return Enumerable.Empty<Statement>();
            }
        }
        #endregion Public Methods

        #region Private Methods
        private IEnumerable<IEnumerable<Statement>> AllPaths(Statement initialStatement, Func<Statement, IEnumerable<Statement>> getNextStatements)
        {
            // Start with an empty path and the statements after the initial statement.
            var path = new List<Statement>();
            var nextStatementsEnumerator = new Stack<IEnumerator<Statement>>();
            nextStatementsEnumerator.Push(getNextStatements(initialStatement).GetEnumerator());

            while (nextStatementsEnumerator.Count > 0)
            {
                // Get the statements we are currently working on.
                IEnumerator<Statement> nextStatements = nextStatementsEnumerator.Peek();

                if (nextStatements.MoveNext())
                {
                    var nextStatement = nextStatements.Current;
                    if (path.Contains(nextStatement))
                    {
                        // We looped back upon ourselves so return the current path here.
                        yield return path.ToArray();
                    }
                    else
                    {
                        // Add to the path we are currently walking.
                        path.Add(nextStatement);

                        // Get the next statements, if any.
                        var enumerable = getNextStatements(nextStatement);
                        if (enumerable.Any())
                        {
                            // There are next statements so add them to the list to process.
                            nextStatementsEnumerator.Push(enumerable.GetEnumerator());
                        }
                        else
                        {
                            // This is the end of the path so return the path here.
                            yield return path.ToArray();
                            path.RemoveAt(path.Count - 1);
                        }
                    }
                }
                else
                {
                    // We are done looking at these next statements so pop and move back a level.
                    nextStatementsEnumerator.Pop();
                    if (path.Count > 0)
                    {
                        path.RemoveAt(path.Count - 1);
                    }
                }
            }
        }

        private IEnumerable<Statement> AllStatements(Statement statement, Func<Statement, IEnumerable<Statement>> nextStatements)
        {
            HashSet<Statement> visitedStatements = new HashSet<Statement>();
            Stack<Statement> statementsToProcess = new Stack<Statement>();
            statementsToProcess.Push(statement);

            while (statementsToProcess.Count > 0)
            {
                // Get the next statement to process. Skip it if we've already processed it.
                Statement currentStatement = statementsToProcess.Pop();
                foreach (Statement childStatement in nextStatements(currentStatement))
                {
                    // Only process statements that haven't already been seen.
                    if (!visitedStatements.Contains(childStatement))
                    {
                        // Return this statement and add it to the list that have already been processed.
                        yield return childStatement;
                        visitedStatements.Add(childStatement);

                        // Add this as another statement to follow.
                        statementsToProcess.Push(childStatement);
                    }
                }
            }
        }

        private void BuildStatementFlow()
        {
            this.statements = new Dictionary<Statement, StatementInfo>();

            this.ProcessAllStatements();
        }

        private void EnsureFlowHasBeenBuilt()
        {
            if (this.statements == null)
            {
                this.BuildStatementFlow();
            }
        }

        private void ProcessAllStatements()
        {
            var branches = new Stack<BranchInfo>();
            var firstStatements = new Dictionary<Block, Statement>();

            StatementInfo previousStatementInfo = null;
            List<Block> blocksToGetFirstStatementFor = new List<Block>();
            foreach (Statement statement in this.Method.Body.ExpandBlocks(true).IgnoreNops())
            {
                if (statement.NodeType == NodeType.Block)
                {
                    blocksToGetFirstStatementFor.Add((Block)statement);
                }
                else
                {
                    // If we are looking for the first statement of any block then update it here.
                    if (blocksToGetFirstStatementFor.Count > 0)
                    {
                        foreach (Block block in blocksToGetFirstStatementFor)
                        {
                            firstStatements.Add(block, statement);
                        }

                        blocksToGetFirstStatementFor.Clear();
                    }

                    // Add this statement to the collection.
                    StatementInfo statementInfo = new StatementInfo(statement);
                    this.statements.Add(statement, statementInfo);

                    // If there is a previous statement then add it to the list previous statements and add the current
                    // statement to its following statements.
                    if (previousStatementInfo != null)
                    {
                        statementInfo.AddPrevious(previousStatementInfo.Statement);
                        previousStatementInfo.AddFollowing(statement);
                    }

                    if (statement.NodeType == NodeType.Branch)
                    {
                        Branch branch = (Branch)statement;

                        if (branch.Condition == null)
                        {
                            // If this isn't a conditional branch then we don't fall through to the next statement.
                            // That means we stop walking this code path.
                            previousStatementInfo = null;
                            continue;
                        }
                        else
                        {
                            branches.Push(new BranchInfo(statementInfo, branch.Target));
                        }
                    }
                    else if (statement.NodeType == NodeType.Return)
                    {
                        // Nothing follows a return statement.
                        previousStatementInfo = null;
                        continue;
                    }

                    // We are moving to the next statement so the current statement is now the previous one.
                    previousStatementInfo = statementInfo;
                }
            }

            this.ProcessBranches(branches, firstStatements);
        }

        private void ProcessBranches(Stack<BranchInfo> branches, Dictionary<Block, Statement> firstStatements)
        {
            foreach (BranchInfo branchInfo in branches)
            {
                Statement blocksFirstStatement = firstStatements[branchInfo.TargetBlock];
                StatementInfo targetStatementInfo = this.statements[blocksFirstStatement];
                targetStatementInfo.AddPrevious(branchInfo.PreviousStatement.Statement);
                branchInfo.PreviousStatement.AddFollowing(targetStatementInfo.Statement);
            }
        }
        #endregion Private Methods

        #region Private Classes
        /// <summary>
        /// Holds information about a <see cref="Branch"/>.
        /// </summary>
        private class BranchInfo
        {
            private readonly Block targetBlock;
            private readonly StatementInfo previousStatement;

            #region Constructors
            /// <summary>
            /// Initializes a new instance of the <see cref="BranchInfo"/> class.
            /// </summary>
            /// <param name="previousStatement">The previous statement.</param>
            /// <param name="targetBlock">The target block.</param>
            public BranchInfo(StatementInfo previousStatement, Block targetBlock)
            {
                this.previousStatement = previousStatement;
                this.targetBlock = targetBlock;
            }
            #endregion Constructors

            #region Public Properties
            /// <summary>
            /// Gets the previous statement.
            /// </summary>
            /// <value>The previous statement.</value>
            public StatementInfo PreviousStatement
            {
                get { return this.previousStatement; }
            }

            /// <summary>
            /// Gets the target block.
            /// </summary>
            /// <value>The target block.</value>
            public Block TargetBlock
            {
                get { return this.targetBlock; }
            }
            #endregion Public Properties
        }
        #endregion Private Classes
    }
}